from pydantic import BaseModel, Field


class POSCARRequest(BaseModel):
    """POSCAR文件标准化处理请求模型"""

    poscar_content: str = Field(..., description="POSCAR文件的完整内容，包含晶格参数、原子种类、原子位置等信息")


class TQCDataRequest(POSCARRequest):
    """拓扑量子化学(TQC)数据分析请求模型"""

    tqc_data: str = Field(..., description="TQC数据文件内容，包含占据带信息。每行代表一个k点，1表示占据，0表示未占据")


class CRRequest(BaseModel):
    """兼容性关系检查请求模型"""

    poscar_content: str | None = Field(
        None, description="可以不传，用来定位cache目录，如果传了，则cache目录为poscar_content的hash值"
    )
    tqc_data: str = Field(..., description="TQC数据文件内容，包含占据带信息。每行代表一个k点，1表示占据，0表示未占据")
    og_number: int | None = Field(
        None, description="OG磁性空间群编号(可选)。范围：1-1651。如果不提供，将自动确定", gt=0, le=1651
    )


class MSGRequest(POSCARRequest):
    """磁性空间群转换请求模型"""

    sg_number: int = Field(..., description="标准空间群编号，对应国际晶体学表。范围：1-230", gt=0, le=230)
    og_number: int = Field(
        ..., description="OG磁性空间群编号，Opechowski-Guccione编号系统。范围：1-1651", gt=0, le=1651
    )


class MOM2MSGRequest(POSCARRequest):
    """从带磁矩的POSCAR生成MSG相关文件请求模型"""

    pass


class SIRequest(CRRequest):
    """自旋指标计算请求模型"""

    pass


# 响应模型
class POSCARResponse(BaseModel):
    """POSCAR处理响应模型"""

    status: str = Field(description="处理状态，通常为 'success'")
    poscar_std: str = Field(description="标准化后的POSCAR文件内容")
    abrs: str = Field(description="ABR输出内容，包含原子轨道信息")

    model_config = {
        "json_schema_extra": {
            "example": {
                "status": "success",
                "poscar_std": "Ti2C\\n1.0\\n3.068 0.000 0.000\\n0.000 3.068 0.000\\n0.000 0.000 20.000\\nTi C\\n2 1\\nDirect\\n0.000 0.000 0.500\\n0.667 0.333 0.500\\n0.333 0.667 0.500",
                "abrs": "2a@1a_1: Ti(0,0,1/2): dxy dxz dyz dx2-y2 dz2\\n1b@1b_1: C(1/3,2/3,1/2): s px py pz",
            }
        }
    }


class EBRResponse(BaseModel):
    """eBR分解响应模型"""

    status: str = Field(description="处理状态，通常为 'success'")
    soc: int = Field(description="是否包含自旋轨道耦合，0=无SOC，1=有SOC")
    abrs: str = Field(description="ABR输出内容")
    num_solutions: int = Field(description="eBR分解解的数量，0表示拓扑材料")
    decompose_ebr_content: str = Field(description="eBR分解详细结果")
    is_topological: bool = Field(description="是否为拓扑材料")
    conclusion: str = Field(description="分析结论")
    additional_data: dict | None = Field(None, description="附加数据，包含额外的详细信息或说明")

    model_config = {
        "json_schema_extra": {
            "example": {
                "status": "success",
                "soc": 0,
                "abrs": "2a@1a_1: Ti(0,0,1/2): dxy dxz dyz dx2-y2 dz2\\n1b@1b_1: C(1/3,2/3,1/2): s px py pz",
                "num_solutions": 0,
                "decompose_ebr_content": "0\\n(eBR分解结果)",
                "is_topological": True,
                "conclusion": "It's topological, since there is no solution for eBR decomposition. Please further check the compatibility relations (CRs).",
                "additional_data": {
                    "details": "If satisfying CRs, it is a topological insulator; otherwise, it is a topological semimetal."
                },
            }
        }
    }


class ABRResponse(BaseModel):
    """aBR分解响应模型"""

    status: str = Field(description="处理状态，通常为 'success'")
    soc: int = Field(description="是否包含自旋轨道耦合，0=无SOC，1=有SOC")
    abrs: str = Field(description="abrs输出内容")
    num_ebr_solutions: int = Field(description="eBR分解解的数量")
    num_abr_solutions: int = Field(description="aBR分解解的数量")
    decompose_abr_content: str = Field(description="aBR分解详细结果")
    type: str = Field(description="材料类型：AI(原子绝缘体)、OAI1/OAI2(阻塞原子绝缘体)、TI(拓扑绝缘体)")
    conclusion: str = Field(description="拓扑性质结论")
    additional_data: dict | None = Field(None, description="附加数据，包含额外的详细信息或说明")

    model_config = {
        "json_schema_extra": {
            "example": {
                "status": "success",
                "soc": 0,
                "abrs": "2a@1a_1: Ti(0,0,1/2): dxy dxz dyz dx2-y2 dz2\\n1b@1b_1: C(1/3,2/3,1/2): s px py pz",
                "num_ebr_solutions": 0,
                "num_abr_solutions": 0,
                "decompose_abr_content": "0\\n(aBR分解结果)",
                "type": "TI",
                "conclusion": "Topological insulator",
                "additional_data": {
                    "details": "Please further check the compatibility relations (CRs). If satisfying CRs, it is a topological insulator; otherwise, it is a topological semimetal."
                },
            }
        }
    }


class CRResponse(BaseModel):
    """兼容性关系响应模型"""

    status: str = Field(description="处理状态，通常为 'success'")
    soc: int = Field(description="是否包含自旋轨道耦合，0=无SOC，1=有SOC")
    cr_result: list[str] = Field(description="不满足兼容性关系的对称线或k点列表，空列表表示满足所有兼容性关系")
    conclusion: str = Field(description="兼容性关系检查结论")

    model_config = {
        "json_schema_extra": {"example": {"status": "success", "soc": 0, "cr_result": [], "conclusion": "Satisfy CR"}}
    }


class MSGResponse(BaseModel):
    """磁性空间群转换响应模型"""

    status: str = Field(description="处理状态，通常为 'success'")
    sg_number: int = Field(description="标准空间群编号")
    og_number: int = Field(description="OG磁性空间群编号")
    msg_type: int = Field(description="磁性空间群类型")
    poscar_msg: str = Field(description="磁性POSCAR文件内容")
    incar: str = Field(description="INCAR文件内容，包含磁性计算参数")
    kpoints: str = Field(description="KPOINTS文件内容")

    model_config = {
        "json_schema_extra": {
            "example": {
                "status": "success",
                "sg_number": 225,
                "og_number": 1435,
                "msg_type": 4,
                "poscar_msg": "Ti2C\\n1.0\\n3.068 0.000 0.000\\n...",
                "incar": "LSORBIT = T\\nLNONCOLLINEAR = T\\nSAXIS = 0 0 1\\nISTART = 0\\nICHARG = 2",
                "kpoints": "Automatic mesh\\n0\\nMonkhorst\\n8 8 2\\n0. 0. 0.",
                "magvec": "0.0 0.0 4.0\\n0.0 0.0 -4.0",
            }
        }
    }


class SIResponse(BaseModel):
    """自旋指标响应模型"""

    status: str = Field(description="处理状态，通常为 'success'")
    soc: int = Field(description="是否包含自旋轨道耦合，0=无SOC，1=有SOC")
    og_number: int = Field(description="OG磁性空间群编号")
    si_result: str = Field(description="自旋指标计算结果")

    model_config = {
        "json_schema_extra": {
            "example": {
                "status": "success",
                "og_number": 1435,
                "soc": 1,
                "si_result": "SI_x = 0, SI_y = 1, SI_z = 0\\n磁性拓扑绝缘体",
            }
        }
    }
